#!/usr/bin/env python2.7

# layer_tracker.py

import os, sys
import signal
#import re
#from datetime import datetime, timedelta
#from dateutil.parser import parse
import numpy as np
from scipy.ndimage import gaussian_filter1d
import matplotlib
matplotlib.use('Qt4Agg')
matplotlib.rcParams['backend.qt4']='PySide'
import matplotlib.pyplot as plt
from matplotlib.colors import colorConverter
from matplotlib.backends.backend_qt4agg import FigureCanvasQTAgg as FigureCanvas
from matplotlib.backends.backend_qt4agg import NavigationToolbar2QTAgg as NavigationToolbar
from matplotlib.figure import Figure
#from mpl_toolkits.basemap import Basemap
#from matplotlib import mlab
#from netCDF4 import Dataset
from PySide.QtCore import *
from PySide.QtGui  import *
from ui_mlh import *
from averager import *
from lidarimage import *
from bender import *
from grad import *
from datelist import *
from bumper import *
from peaker import *
from filer import *

from matplotlib.path import Path
from matplotlib.widgets import *
#from mpl_toolkits.basemap import Basemap
#from matplotlib import mlab
from scipy.ndimage import gaussian_filter1d
from netCDF4 import Dataset
from metlib.kits import *
from metlib.color.color_collections import *
from metlib.lidar import *
from metlib.plot.artist_util import ArtistIDManager
from lidarmess.lidar_extra import *

sin75 = np.sin(np.deg2rad(75))

cmapd = {'jet':plt.cm.jet}


bump_setting_attrs = [
            'averager.aver_min', 
            'averager.on',
            'averager.strict_decloud',
            'averager.n_down',
            'averager.n_spread',
            'bender.hmin',
            'bender.h',
            'bender.xmin',
            'bender.on',
            'gradder.gradtype',
            'gradder.sigma',
            'bumper.fall_v0',
            'bumper.fall_a',
            'bumper.intensity_thres',
            'bumper.intensity_thres2',
            'bumper.intensity_n',
            'bumper.window_len',
            'bumper.hthres_upper',
            'bumper.hthres_lower',
            ]

gui_setting_attrs = [
            'averager.aver_min', 
            'averager.on',
            'averager.strict_decloud',
            'averager.n_down',
            'averager.n_spread',
            'bender.hmin',
            'bender.h',
            'bender.xmin',
            'bender.on',
            'gradder.gradtype',
            'gradder.sigma',
            'peaker.pos_offset',
            'peaker.pos_ratio',
            'peaker.pos_thres',
            'peaker.neg_offset',
            'peaker.neg_ratio',
            'peaker.neg_thres',
            'bumper.fall_v0',
            'bumper.fall_a',
            'bumper.intensity_thres',
            'bumper.intensity_thres2',
            'bumper.intensity_n',
            'bumper.window_len',
            'bumper.hthres_upper',
            'bumper.hthres_lower',
            'rcsimage.cmapname',
            'rcsimage.vmin',
            'rcsimage.vmax',
            'gradimage.cmapname',
            'gradimage.vmin',
            'gradimage.vmax',
            'filer.settingfname',
            'filer.dirname',
            'filer.prefixrcs',
            'filer.prefixgrad',
            'filer.prefixpeak',
            'filer.subdir'
            ]


def guiquit(*args):
    print "Get INT!!!"
    QApplication.quit()
    sys.exit()

def curve2(x):
    return 1.0 - (x-1.0) ** 2

def toQColor(color):
    return QColor.fromRgbF(*colorConverter.to_rgba(color))

class MLHGUI(QMainWindow, Ui_Form, Setter):
    def __init__(self, parent=None, rc_file=''):
        super(MLHGUI, self).__init__(parent)
        self.setupUi(self)
        
        self.rcsimage = LidarImage(parent=self)
        self.rcsimageLayout.addWidget(self.rcsimage)
        self.gradimage = LidarImage(parent=self)
        self.gradimageLayout.addWidget(self.gradimage)
        self.peakimage = LidarPeakImage(parent=self, ref_image=self.rcsimage)
        self.peakimageLayout.addWidget(self.peakimage)

        self.gradimage.cmapname= 'cm_GateBWR2'
        self.gradimage.vmin = -0.05
        self.gradimage.vmax = 0.05

        self.averager = Averager(parent=self)
        self.toolboxLayout.addWidget(self.averager)
        self.averager.applyButton.clicked.connect(self.on_averager_change)
        self.averager.checkBox.stateChanged.connect(self.on_averager_change)

        self.bender = Bender(parent=self)
        self.toolboxLayout.addWidget(self.bender)
        self.bender.applyButton.clicked.connect(self.on_bender_change)
        self.bender.checkBox.stateChanged.connect(self.on_bender_change)

        self.gradder = Grad(parent=self)
        self.toolboxLayout.addWidget(self.gradder)

        self.peaker = Peaker(parent=self)
        self.toolboxLayout.addWidget(self.peaker)
        self.peaker.applyButton.clicked.connect(self.on_peak_ms_change)

        self.bumper = Bumper(parent=self)
        self.toolboxLayout.addWidget(self.bumper)
        self.bumper.bumpButton.clicked.connect(self.bump)
        self.bumper.remove_Button.clicked.connect(self.remove_bumpline)
        self.bumper.remove_all_Button.clicked.connect(self.peakimage.clear_extra_data)
        self.bumper.ltrim_Button.clicked.connect(self.trim_bumpline_l)
        self.bumper.rtrim_Button.clicked.connect(self.trim_bumpline_r)
        
        self.bumplines = []

        self.dtlist = Datelist(parent=self)
        self.toolboxLayout.addWidget(self.dtlist)
        self.dtlist.datelist.currentRowChanged.connect(self.switch2date)

        self.tabWidget.setCurrentIndex(0)

        self.filer = Filer(parent=self)
        self.toolboxLayout.addWidget(self.filer)
        self.filer.saveButton.clicked.connect(self.save)
        self.filer.savesettingButton.clicked.connect(self.save_all_setting)

        self.colorbox = ColorBox(pop_poster_random)

        self.shortcut_save = QShortcut(QKeySequence("Ctrl+s"), self, self.save)
        self.shortcut_triml = QShortcut(QKeySequence("z"), self, self.trim_bumpline_l)
        self.shortcut_trimr = QShortcut(QKeySequence("x"), self, self.trim_bumpline_r)
        self.shortcut_bump =  QShortcut(QKeySequence("Space"), self, self.bump)
        self.shortcut_remove =  QShortcut(QKeySequence("Del"), self, self.remove_bumpline)
        self.shortcut_removeall =  QShortcut(QKeySequence("Ctrl+Del"), self, self.remove_bumpline)
        self.shortcut_switchrcs =  QShortcut(QKeySequence("1"), self, self.switch_tab_rcs)
        self.shortcut_switchgrad =  QShortcut(QKeySequence("2"), self, self.switch_tab_grad)
        self.shortcut_switchpeak =  QShortcut(QKeySequence("3"), self, self.switch_tab_peak)

        if rc_file:
            self.load_setting(rc_file)

    def load_dataset(self, lidardir, fname_pattern=r'\d{10}.*\.nc'):
        self.display("Loading dataset: %s" % lidardir)
        self.lidardir = lidardir
        lidarfnames = list_all_file(lidardir, fname_pattern=fname_pattern)
        lidardtstrs = [os.path.basename(fn)[0:8] for fn in lidarfnames]
        self.file_i = limited_int(0, 0, len(lidardtstrs)-1)
        
        self.dataset = []
        for fn in lidarfnames:
            self.dataset.append((fn, ))

        self.dtlist.load(lidardtstrs)

        
    def load_file_i(self, i, force=False):
        old_i = int(self.file_i)
        if i is None or i is Null:
            try:
                new_i = int(raw_input(
                    '  Input file_i <min:%d, max:%d>: ' % (self.file_i.vmin, self.file_i.vmax)))
                self.file_i.set_value(new_i)
            except Exception as e:
                print e
                return
        elif isinstance(i, (str, unicode)):
            if i == 'next':
                self.file_i += 1
            elif i == 'prev':
                self.file_i -= 1
            elif i.startswith('jump_'):
                jump_i = int(i[5:])
                self.file_i += jump_i
        elif isinstance(i, (int, long, np.integer)):
            self.file_i.set_value(i)
        if force or int(self.file_i) != old_i:
            print "Loading File No. ", self.file_i
            self.dtlist.index = self.file_i

    def switch2date(self, i):
        self.file_i.set_value(int(i))
        self.load_data(i)
        
    def load_data(self, i):
        self.display("Loading data")
        self.lidard_bak = LidarDataset(self.dataset[i])
        begdtstr = self.lidard_bak['datetime'][0].strftime('%Y%m%d')
        enddtstr = self.lidard_bak['datetime'][-1].strftime('%Y%m%d')
        if begdtstr == enddtstr:
            self.filer.fname = begdtstr
        else:
            self.filer.fname = '%s-%s' % (begdtstr, enddtstr)
        # decloud
        self.cbi = self.averager.decloud(self.lidard_bak)

        self.lidare = lidar_extra(self.lidard_bak)
        self.max_hi = height_to_index(5000.0, self.lidard_bak, self.lidare['elev_angle'])
        self.lidard_bak.resize_bin(0, self.max_hi)
        if 'energy corrected' not in self.lidard_bak['desc']:
            correct_energy(self.lidard_bak)
        fill_lower_part(self.lidard_bak, self.lidare['valid_index'], 1)

        self.rcsimage.vmin = self.lidare['norm_vmin']
        self.rcsimage.vmax = self.lidare['norm_vmax']
        
        self.chn = self.lidare['chn']
        
        self.apply_basicdata()
        self.apply_bend()
        self.apply_grad()
        self.apply_peak()
        self.update_axs()
        
    def update_axs(self):
        pass
    
    def apply_basicdata(self):
        self.display("Applying basic data ...")
        self.lidard = self.lidard_bak.copy()
        self.averager.average(self.lidard)
        self.dts = self.lidard['datetime']
        self.heights = self.lidard['distance'] * self.lidare['sin_elev']
        self.norm = self.lidard['data'][:,self.chn,:]
        self.rcs  = np.log10(self.norm)

        self.rcsimage.load_data(self.dts, self.heights, self.rcs)
        self.rcsimage.load_extra_data('plot', self.dts, self.lidard['cloud_base_height'], 'ko', ms=2.0, label='Cloud Base')
        self.rcsimage.update_fig()


    def apply_grad(self):
        self.display("Applying grad ...")
#        forgrad = self.bendednorm.copy()
#        cbi = np.argmax(~(forgrad > 0.0), axis=-1)
#        print cbi
#        level0_bad = ~(forgrad[:, self.lidare['valid_index']] > 0.0) 
#        print level0_bad
#        for i in range(len(cbi)):
#            if cbi[i] > 0 :
#                print i, cbi[i]
#                forgrad[i, cbi[i]:] = forgrad[i, cbi[i]-1]
#            elif (cbi[i] == 0 and level0_bad[i]):
#                forgrad[i, :] = np.nan
#        self.rcsimage.load_data(self.dts, self.heights, np.log10(forgrad))
#        self.rcsimage.update_fig()
        self.grad = self.gradder.grad(self.bendednorm, self.lidard['bin_size'], fill=True)
        self.gradimage.load_data(self.dts, self.heights, self.grad)
        self.gradimage.update_fig()

    def apply_bend(self):
        self.display("Applying bender ...")
        self.bendednorm = self.bender.bend(self.norm, self.heights)

    def apply_peak(self):
        self.display("Applying peak stuff ...")
        self.peaks_list = measure_peak(self.grad.transpose())
        self.peakimage.load_data(*self.peaker.peaks_to_scatter(self.dts, self.heights, self.peaks_list, pos_color='r', neg_color='b'))
        self.peakimage.clear_extra_data()
        self.bumper.bumpList.clear()
        del self.bumplines[0:len(self.bumplines)]
        self.peakimage.update_fig()

    def on_bender_change(self):
        self.apply_bend()
        self.apply_grad()
        self.apply_peak()

    def on_averager_change(self):
        self.apply_basicdata()
        self.apply_bend()
        self.apply_grad()
        self.apply_peak()

    def on_peak_change(self):
        self.apply_peak()

    def on_peak_ms_change(self):
        self.peakimage.load_data(*self.peaker.peaks_to_scatter(self.dts, self.heights, self.peaks_list, pos_color='r', neg_color='b'))
        self.peakimage.update_fig()

    def bump(self):
        self.display("Bumping")
        dt = self.peakimage.sel_t
        h = self.peakimage.sel_h
        if dt is None or h is None:
            return
        start_i = np.where(self.dts == dt)[0][0]
        time_inds, height_series, intensity_series = self.bumper.bump(self.peaks_list, self.heights, start_i, start_hmin=h*self.bumper.hthres_lower, start_hmax=h*self.bumper.hthres_upper)
        time_series = self.dts[time_inds]
        color = self.colorbox()
        self.peakimage.load_extra_data('plot', time_series, height_series, 'o-', color=color, ms=2.5, lw=1.5, alpha=0.7, picker=5)
        bumpitem = QListWidgetItem()
        bumpitem.setBackground(QBrush(toQColor(color)))
        self.bumper.bumpList.addItem(bumpitem)
        self.bumplines.append([self.get_all_setting(), time_series, height_series])
#        self.peakimage.update_fig()

    def remove_bumpline(self):
        self.display("Removing bumpline # %d" % self.bumper.index)
        self.peakimage.clear_extra_data(self.bumper.index)
        del self.bumplines[self.bumper.index]
        self.bumper.bumpList.takeItem(self.bumper.index)

    def trim_bumpline(self, lr):
        if self.peakimage.sel_line is not None and self.peakimage.sel_line_i is not None:
            i_trim = self.peakimage.sel_line_i
            line_id = self.peakimage.extra_artist_ids[self.peakimage.sel_line]
            extra_i = self.peakimage.index2id.index(line_id)
            dts, heights = self.peakimage.extra_data[extra_i][1:3] 
            if lr == 'l':
                dts = dts[i_trim:]
                heights = heights[i_trim:]
            else:
                dts = dts[:i_trim]
                heights = heights[:i_trim]
            self.peakimage.extra_data[extra_i][1] = dts
            self.peakimage.extra_data[extra_i][2] = heights
            self.peakimage.sel_line.set_data(dts, heights)
            self.peakimage.canvas.draw()
            self.bumplines[extra_i][1] = dts
            self.bumplines[extra_i][2] = heights
       
    def trim_bumpline_l(self):
        self.display("< Trimming")
        self.trim_bumpline('l')

    def trim_bumpline_r(self):
        self.display("Trimming >")
        self.trim_bumpline('r')

    def save(self):
        self.display( "Saving " + self.filer.fname )
        sep = '/' if self.filer.subdir else '_'
        rcsfname = expand_path('%s/%s%s%s.png' % (self.filer.dirname, self.filer.prefixrcs, sep, self.filer.fname))
        gradfname = expand_path('%s/%s%s%s.png' % (self.filer.dirname, self.filer.prefixgrad, sep, self.filer.fname))
        peakfname = expand_path('%s/%s%s%s.png' % (self.filer.dirname, self.filer.prefixpeak, sep, self.filer.fname))
        bumpfname = expand_path('%s/%s%s%s.bump' % (self.filer.dirname, self.filer.prefixbump, sep, self.filer.fname))
        for fname in (rcsfname, gradfname, peakfname, bumpfname):
            dirname = os.path.dirname(fname)
            force_makedirs(dirname)

        self.rcsimage.fig.savefig(rcsfname)
        self.gradimage.fig.savefig(gradfname)
        self.peakimage.fig.savefig(peakfname)

        with open(bumpfname, 'w') as bumpfile:
            bumpfile.write('datetime,height\n')
            for setting, dts, heights in self.bumplines:
                for setting_item in setting:
                    bumpfile.write('# ')
                    bumpfile.write(setting_item)
                    bumpfile.write('\n')
                for dt, h in zip(dts, heights):
                    bumpfile.write('%s,%.2f\n' % (dt, h))



    def get_all_setting(self, prefix=''):
        return self.get_setting(attrs=bump_setting_attrs, prefix=prefix
            )

    def save_all_setting(self):
        self.save_setting(self.filer.settingfname, attrs=gui_setting_attrs)

    def display(self, text):
        self.label.setText(text)
        QTimer.singleShot(3000, lambda: self.label.setText(""))

    def switch_tab(self, n):
        self.tabWidget.setCurrentIndex(n)

    def switch_tab_rcs(self):
        self.switch_tab(0)

    def switch_tab_grad(self):
        self.switch_tab(1)

    def switch_tab_peak(self):
        self.switch_tab(2)

if __name__ == '__main__':
    signal.signal(signal.SIGINT, guiquit)

    reset = False
#    get_sys_argv(['detaildir', 'lidardir', 'station'])
#    get_sys_argv(['data_dir', 'image_dir', 'out_dir', 'station', ('resolution', int), ('beg_date', T), ('end_date', T)], [], [('reset', parse_bool)] )
#    force_makedirs(out_dir)
#    fnames_list = []
#    for dt in datetime_range(beg_date, end_date, '1d'):
#        dtstr = dt.strftime('%Y%m%d')
#        data_fn = '%s/%s.txt' % (data_dir, dtstr)
#        image_fn = '%s/%s.png' % (image_dir, dtstr)
#        out_fn = '%s/%s.txt' % (out_dir, dtstr)
#        if os.path.exists(data_fn) and os.path.exists(image_fn):
#            fnames_list.append((data_fn, image_fn, out_fn))
    rc = ''
    get_sys_argv(['lidardir'], [], [('rc', str)])
    
    app = QApplication(sys.argv)
    frame = MLHGUI(rc_file=rc)
#    frame.load_dataset(detaildir, lidardir, station)
    frame.load_dataset(lidardir)
    frame.show()
    timer = QTimer()
    timer.start(1000)
    timer.timeout.connect(lambda: None)
    sys.exit(app.exec_())
